/*
 * Copyright © 2024-2025 Wenze Wei
 *
 * This file is part of StadionOS.
 *
 * This Source Code Form is subject to the terms of the GNU General Public License, v. 2.0.
 * You can obtain a copy of the GPL at
 * https://www.gnu.org/licenses/old-licenses/gpl-2.0.html.
 *
 * Software distributed under the License is distributed on an "AS IS" basis,
 * WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License for
 * the specific language governing rights and limitations under the License.
 */

#include "../include/lpi.h"

TNode *create_scalar_node(const char *value) {
    TNode *node = Smalloc(sizeof(TNode));
    node->type = NODE_SCALAR;
    node->data.scalar = strdup(value);
    return node;
}

TNode *create_sequence_node() {
    TNode *node = Smalloc(sizeof(TNode));
    node->type = NODE_SEQUENCE;
    node->data.sequence.items = NULL;
    node->data.sequence.count = 0;
    return node;
}

TNode *create_mapping_node() {
    TNode *node = Smalloc(sizeof(TNode));
    node->type = NODE_MAPPING;
    node->data.mapping.entries = NULL;
    return node;
}

void add_sequence_element(TNode *sequence, TNode *element) {
    sequence->data.sequence.items = Srealloc(
        sequence->data.sequence.items,
        (sequence->data.sequence.count + 1) * sizeof(TNode*)
    );
    sequence->data.sequence.items[sequence->data.sequence.count++] = element;
}

void add_mapping_entry(TNode *mapping, char *key, TNode *value) {
    TMapEntry *entry = Smalloc(sizeof(TMapEntry));
    entry->key = key;
    entry->value = value;
    entry->next = mapping->data.mapping.entries;
    mapping->data.mapping.entries = entry;
}

TNode *parse_yaml(const char *filename) {
    FILE *file = fopen(filename, "rb");
    if (!file) return NULL;

    yaml_parser_t parser;
    if (!yaml_parser_initialize(&parser)) {
        fclose(file);
        return NULL;
    }
    yaml_parser_set_input_file(&parser, file);

    TNode *root = NULL;
    TNode **stack = NULL;
    size_t stack_size = 0;
    char *current_key = NULL;
    int expecting_key = 1;

    yaml_event_t event;
    do {
        if (!yaml_parser_parse(&parser, &event)) break;

        switch (event.type) {
            case YAML_MAPPING_START_EVENT: {
                TNode *node = create_mapping_node();
                if (!root) root = node;
                if (stack_size > 0) {
                    TNode *parent = stack[stack_size - 1];
                    if (parent->type == NODE_SEQUENCE) {
                        add_sequence_element(parent, node);
                    } else if (parent->type == NODE_MAPPING && !expecting_key) {
                        add_mapping_entry(parent, current_key, node);
                        Sfree(current_key);
                        current_key = NULL;
                        expecting_key = 1;
                    }
                }
                stack = Srealloc(stack, ++stack_size * sizeof(TNode *));
                stack[stack_size - 1] = node;
                break;
            }

            case YAML_SEQUENCE_START_EVENT: {
                TNode *node = create_sequence_node();
                if (!root) root = node;
                if (stack_size > 0) {
                    TNode *parent = stack[stack_size - 1];
                    if (parent->type == NODE_SEQUENCE) {
                        add_sequence_element(parent, node);
                    } else if (parent->type == NODE_MAPPING && !expecting_key) {
                        add_mapping_entry(parent, current_key, node);
                        free(current_key);
                        current_key = NULL;
                        expecting_key = 1;
                    }
                }
                stack = realloc(stack, ++stack_size * sizeof(TNode *));
                stack[stack_size - 1] = node;
                break;
            }

            case YAML_SCALAR_EVENT: {
                char *value = strdup((char *)event.data.scalar.value);
                if (stack_size > 0 && stack[stack_size - 1]->type == NODE_MAPPING && expecting_key) {
                    current_key = value;
                    expecting_key = 0;
                } else {
                    TNode *scalar_node = create_scalar_node(value);
                    Sfree(value);
                    if (stack_size > 0) {
                        TNode *parent = stack[stack_size - 1];
                        if (parent->type == NODE_SEQUENCE) {
                            add_sequence_element(parent, scalar_node);
                        } else if (parent->type == NODE_MAPPING && !expecting_key) {
                            add_mapping_entry(parent, current_key, scalar_node);
                            Sfree(current_key);
                            current_key = NULL;
                            expecting_key = 1;
                        }
                    } else {
                        root = scalar_node;
                    }
                }
                break;
            }

            case YAML_MAPPING_END_EVENT:
            case YAML_SEQUENCE_END_EVENT:
                if (stack_size > 0) stack_size--;
                break;

            default:
                break;
        }

        yaml_event_delete(&event);
    } while (event.type != YAML_STREAM_END_EVENT);

    yaml_parser_delete(&parser);
    fclose(file);
    Sfree(stack);
    if (current_key) Sfree(current_key);
    return root;
}

int parse_key_path(const char *key_path, TPathPart **parts, size_t *count) {
    const char *p = key_path;
    TPathPart *temp = NULL;
    size_t temp_size = 0;

    while (*p) {
        if (*p == '[') {
            p++;
            char *end;
            size_t index = strtoul(p, &end, 10);
            if (*end != ']') {
                for (size_t i = 0; i < temp_size; i++) {
                    if (temp[i].type == PATH_KEY) Sfree(temp[i].value.key);
                }
                Sfree(temp);
                return 0;
            }
            p = end + 1;

            TPathPart part = {PATH_INDEX, .value.index = index};
            temp = Srealloc(temp, (temp_size + 1) * sizeof(TPathPart));
            temp[temp_size++] = part;
        } else if (*p == '.') {
            p++;
        } else {
            const char *start = p;
            while (*p && *p != '.' && *p != '[') p++;
            size_t len = p - start;
            char *key = Smalloc(len + 1);
            Sstrncpy(key, start, len);
            key[len] = '\0';

            TPathPart part = {PATH_KEY, .value.key = key};
            temp = Srealloc(temp, (temp_size + 1) * sizeof(TPathPart));
            temp[temp_size++] = part;
        }
    }

    *parts = temp;
    *count = temp_size;
    return 1;
}

const char *find_node_value(TNode *root, const TPathPart *parts, size_t num_parts) {
    TNode *current = root;
    for (size_t i = 0; i < num_parts; i++) {
        if (!current) break;
        TPathPart part = parts[i];

        switch (current->type) {
            case NODE_MAPPING:
                if (part.type != PATH_KEY) return NULL;
                TMapEntry *entry = current->data.mapping.entries;
                while (entry && Sstrcmp(entry->key, part.value.key)) entry = entry->next;
                current = entry ? entry->value : NULL;
                break;

            case NODE_SEQUENCE:
                if (part.type != PATH_INDEX || part.value.index >= current->data.sequence.count)
                    return NULL;
                current = current->data.sequence.items[part.value.index];
                break;

            default:
                return NULL;
        }
    }

    return (current && current->type == NODE_SCALAR) ? current->data.scalar : NULL;
}

void free_node(TNode *node) {
    if (!node) return;
    switch (node->type) {
        case NODE_SCALAR:
            Sfree(node->data.scalar);
            break;
        case NODE_SEQUENCE:
            for (size_t i = 0; i < node->data.sequence.count; i++)
                free_node(node->data.sequence.items[i]);
            Sfree(node->data.sequence.items);
            break;
        case NODE_MAPPING:
            TMapEntry *entry = node->data.mapping.entries;
            while (entry) {
                TMapEntry *next = entry->next;
                Sfree(entry->key);
                free_node(entry->value);
                Sfree(entry);
                entry = next;
            }
            break;
    }
    free(node);
}

/**
 * @brief Extracts a value from a YAML file based on a key path.
 *
 * This function parses the specified YAML file and retrieves the value associated with the given
 * key path. The key path supports nested mappings (objects) and sequences (arrays), using dot
 * notation (`.`) for nested fields and square brackets (`[]`) for array indices.
 *
 * Example key paths:
 * - "person.name": Extracts the `name` field under the `person` object.
 * - "person.address[0].city": Extracts the `city` field of the first element in the `address` array.
 *
 * @param file The path to the YAML file. Must be a valid file path.
 * @param key  The key path specifying the value to extract. The key path uses dot-separated field
 *             names and array indices (e.g., "person.hobbies[1].type").
 *
 * @return Returns a dynamically allocated string representation of the value associated with the
 *         key path. The caller is responsible for freeing the returned string using `free()`.
 *         If the value is not found or an error occurs, returns `NULL`.
 *
 * @note
 * - Array indices are zero-based.
 * - If the YAML file is malformed or the key path is invalid, the function returns `NULL`.
 * - The returned string is dynamically allocated and must be freed by the caller.
 *
 * @example
 * The following example demonstrates how to use the `Syaml()` function:
 * ```c
 * char *value = Syaml("test.yaml", "person.address[0].city");
 * if (value) {
 *     Sprint("City: %s\n", value);
 *     Sfree(value); // Free the dynamically allocated memory
 * } else {
 *     Sprint("Value not found or error occurred.\n");
 * }
 * ```
 */
char* Syaml(const char *file, const char *key) {
    TNode *root = parse_yaml(file);
    if (!root) return NULL;

    TPathPart *parts = NULL;
    size_t num_parts;
    if (!parse_key_path(key, &parts, &num_parts)) {
        free_node(root);
        return NULL;
    }

    const char *value = find_node_value(root, parts, num_parts);
    char *result = value ? strdup(value) : NULL;

    for (size_t i = 0; i < num_parts; i++)
        if (parts[i].type == PATH_KEY) Sfree(parts[i].value.key);
    Sfree(parts);
    free_node(root);

    return result;
}
